// 🚀 高性能文件传输客户端 - Rust实现
// 支持3MB分片传输 + MD5校验 + 30协程并发 + 零拷贝优化
use std::env;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use md5;
use crossbeam::channel::{bounded, unbounded, Sender, Receiver};
use std::time::{SystemTime, UNIX_EPOCH};
use std::path::Path;
use glob::glob;
use indicatif::{ProgressBar, ProgressStyle};

// 📦 常量定义
const CHUNK_SIZE: usize = 3 * 1024 * 1024;  // 3MB分片大小
const CONCURRENT_WORKERS: usize = 30;       // 并发工作协程数

// 📤 分片信息结构体
#[derive(Debug, Clone)]
struct ChunkInfo {
    index: usize,      // 分片索引
    offset: u64,       // 文件偏移量
    size: usize,       // 分片大小
}

// 📋 传输任务结构体
#[derive(Debug, Clone)]
struct TransferTask {
    chunk_info: ChunkInfo,
    file_path: String,
    server_addr: String,
    task_id: String,   // 添加任务ID
}

// 🚀 主函数入口
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    // 📥 获取命令行参数
    let args: Vec<String> = env::args().collect();
    if args.len() < 3 {
        eprintln!("❌ 使用方法:");
        eprintln!("  单文件传输: {} <服务器地址> <本地文件路径> <远程保存路径>", args[0]);
        eprintln!("  多文件传输: {} <服务器地址> <本地文件路径1> [本地文件路径2] [...]", args[0]);
        eprintln!("  默认端口为8888，如需指定其他端口请使用 服务器地址:端口 的格式");
        eprintln!("  支持通配符，例如: {} <服务器地址> *.mp4", args[0]);
        std::process::exit(1);
    }

    // 解析服务器地址，如果没有指定端口则默认使用8888
    let server_addr = if args[1].contains(':') {
        args[1].clone()
    } else {
        format!("{}:{}", args[1], 8888)
    };
    
    // 收集所有文件路径参数
    let file_paths: Vec<String> = args[2..].to_vec();
    
    // 📂 查找匹配的文件
    let mut files = Vec::new();
    for pattern in &file_paths {
        let matched_files = find_files(pattern)?;
        files.extend(matched_files);
    }
    
    if files.is_empty() {
        eprintln!("❌ 未找到匹配的文件");
        std::process::exit(1);
    }

    println!("🚀 启动文件传输客户端...");
    println!("📍 服务器: {}", server_addr);
    println!("📄 找到 {} 个文件", files.len());
    
    // 打印找到的文件列表
    for (index, file_path) in files.iter().enumerate() {
        println!("   {}. {}", index + 1, file_path);
    }

    // 🔁 串行传输所有文件
    for (index, file_path) in files.iter().enumerate() {
        println!("\n📤 传输进度: {}/{}", index + 1, files.len());
        
        // 确定远程文件路径 (多文件传输模式，使用文件名作为远程路径)
        let remote_file_path = Path::new(file_path)
            .file_name()
            .unwrap_or_default()
            .to_string_lossy()
            .to_string();

        // 传输单个文件
        transfer_single_file(&server_addr, file_path, &remote_file_path).await?;
    }

    println!("\n🎉 所有文件传输完成");
    Ok(())
}

// 🔍 查找匹配的文件
fn find_files(pattern: &str) -> Result<Vec<String>, Box<dyn std::error::Error + Send + Sync>> {
    let mut files = Vec::new();
    
    // 检查是否是通配符模式
    if pattern.contains('*') || pattern.contains('?') || pattern.contains('[') {
        // 使用glob匹配
        for entry in glob(pattern)? {
            match entry {
                Ok(path) => {
                    if path.is_file() {
                        files.push(path.to_string_lossy().to_string());
                    }
                }
                Err(e) => eprintln!("❌ glob匹配错误: {}", e),
            }
        }
    } else {
        // 单个文件
        let path = Path::new(pattern);
        if path.is_file() {
            files.push(pattern.to_string());
        }
    }
    
    Ok(files)
}

// 📤 传输单个文件
async fn transfer_single_file(
    server_addr: &str,
    local_file_path: &str,
    remote_file_path: &str,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    println!("📄 本地文件: {}", local_file_path);
    println!("💾 远程路径: {}", remote_file_path);

    // 🔍 获取文件信息
    let file = File::open(local_file_path)?;
    let metadata = file.metadata()?;
    let file_size = metadata.len();
    let total_chunks = ((file_size as f64) / (CHUNK_SIZE as f64)).ceil() as usize;
    
    // 🆔 生成任务ID
    let task_id = generate_task_id();
    
    println!("📊 文件大小: {} bytes ({:.2} MB)", file_size, file_size as f64 / 1024.0 / 1024.0);
    println!("📦 分片数量: {}", total_chunks);
    println!("🆔 任务ID: {}", task_id);

    // 🔐 计算文件MD5 (用于完整性校验)
    let md5_hash = calculate_file_md5(local_file_path)?;
    println!("🔑 文件MD5: {}", md5_hash);

    // 🤝 建立握手连接，发送文件元信息
    let handshake_result = handshake_with_server(
        server_addr, 
        &task_id,
        remote_file_path, 
        file_size, 
        total_chunks, 
        &md5_hash
    ).await?;

    if !handshake_result {
        eprintln!("❌ 服务器握手失败");
        return Err("服务器握手失败".into());
    }
    println!("🤝 与服务器握手成功");

    // ⚡ 启动并发传输
    println!("🏃 启动并发传输 ({}个工作协程)", CONCURRENT_WORKERS);
    concurrent_transfer(
        local_file_path.to_string(),
        server_addr.to_string(),
        task_id.clone(),
        file_size,
        total_chunks,
    ).await?;

    // 📢 通知服务器传输已完成
    println!("📤 通知服务器传输已完成...");
    let completion_result = notify_transfer_complete(server_addr, &task_id).await?;
    
    if completion_result {
        println!("✅ 服务器已确认传输完成");
    } else {
        eprintln!("❌ 服务器未确认传输完成");
        return Err("服务器未确认传输完成".into());
    }

    println!("🎉 文件传输完成");
    Ok(())
}

// 🔐 计算文件MD5值
fn calculate_file_md5(file_path: &str) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
    println!("🔍 正在计算MD5...");
    let mut file = File::open(file_path)?;
    let mut hasher = md5::Context::new();
    let mut buffer = [0u8; 64 * 1024]; // 64KB缓冲区
    
    loop {
        let n = file.read(&mut buffer)?;
        if n == 0 {
            break;
        }
        hasher.consume(&buffer[..n]);
    }
    
    let digest = hasher.compute();
    Ok(format!("{:x}", digest))
}

// 🆔 生成任务ID
fn generate_task_id() -> String {
    let start = SystemTime::now();
    let since_the_epoch = start.duration_since(UNIX_EPOCH)
        .expect("Time went backwards");
    format!("{}", since_the_epoch.as_millis())
}

// 🤝 与服务器握手，发送文件元信息
async fn handshake_with_server(
    server_addr: &str,
    task_id: &str,
    remote_file_path: &str,
    file_size: u64,
    total_chunks: usize,
    md5_hash: &str,
) -> Result<bool, Box<dyn std::error::Error + Send + Sync>> {
    // 📡 建立TCP连接
    let mut stream = TcpStream::connect(server_addr).await?;
    
    // 📦 构造握手消息: "HANDSHAKE|任务ID|文件路径|文件大小|分片总数|MD5"
    let handshake_msg = format!(
        "HANDSHAKE|{}|{}|{}|{}|{}\n",
        task_id, remote_file_path, file_size, total_chunks, md5_hash
    );
    
    // 📤 发送握手消息
    stream.write_all(handshake_msg.as_bytes()).await?;
    
    // 📥 接收服务器响应
    let mut buffer = [0; 1024];
    let n = stream.read(&mut buffer).await?;
    let response = String::from_utf8_lossy(&buffer[..n]);
    
    // ✅ 检查握手结果
    Ok(response.trim() == "HANDSHAKE_OK")
}

// 📢 通知服务器传输已完成
async fn notify_transfer_complete(
    server_addr: &str,
    task_id: &str,
) -> Result<bool, Box<dyn std::error::Error + Send + Sync>> {
    // 📡 建立TCP连接
    let mut stream = TcpStream::connect(server_addr).await?;
    
    // 📦 构造传输完成消息: "TRANSFER_COMPLETE|任务ID"
    let completion_msg = format!("TRANSFER_COMPLETE|{}\n", task_id);
    
    // 📤 发送传输完成消息
    stream.write_all(completion_msg.as_bytes()).await?;
    
    // 📥 接收服务器响应
    let mut buffer = [0; 1024];
    let n = stream.read(&mut buffer).await?;
    let response = String::from_utf8_lossy(&buffer[..n]);
    
    // ✅ 检查服务器响应
    Ok(response.trim() == "COMPLETION_ACK")
}

// ⚡ 并发传输所有分片（零拷贝优化版本）
async fn concurrent_transfer(
    file_path: String,
    server_addr: String,
    task_id: String,
    file_size: u64,
    total_chunks: usize,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    // 📦 创建任务通道（生产者-消费者模型）
    let (sender, receiver) = unbounded::<TransferTask>();
    
    // 📋 生成所有分片任务
    let tasks = generate_transfer_tasks(&file_path, &server_addr, &task_id, file_size, total_chunks);
    
    // 📊 创建进度条
    let pb = ProgressBar::new(total_chunks as u64);
    pb.set_style(ProgressStyle::default_bar()
        .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({percent}%) {msg}")?
        .progress_chars("#>-"));
    
    let progress = std::sync::Arc::new(std::sync::Mutex::new(pb));
    
    // 📤 启动任务生产者
    let sender_clone = sender.clone();
    let progress_clone = progress.clone();
    tokio::task::spawn_blocking(move || {
        for task in tasks {
            if let Err(_) = sender_clone.send(task) {
                break;
            }
        }
        
        // 生产者完成，设置进度条消息
        let progress_lock = progress_clone.lock().unwrap();
        progress_lock.set_message("分片发送完成");
    });

    // 👷 创建工作协程池
    let mut workers = vec![];
    for i in 0..CONCURRENT_WORKERS {
        let receiver_clone = receiver.clone();
        let progress_clone = progress.clone();
        let worker = tokio::spawn(async move {
            if let Err(e) = worker_routine(i, receiver_clone, progress_clone).await {
                eprintln!("❌ 工作协程{}出错: {}", i, e);
            }
        });
        workers.push(worker);
    }

    // 🚫 关闭发送端，确保接收端能正常退出
    drop(sender);

    // 🔄 等待所有工作协程完成
    for worker in workers {
        worker.await;
    }
    
    // 🎉 完成进度条
    let progress_lock = progress.lock().unwrap();
    progress_lock.finish_with_message("所有分片发送完成");

    Ok(())
}

// 📋 生成所有传输任务
fn generate_transfer_tasks(
    file_path: &str,
    server_addr: &str,
    task_id: &str,
    file_size: u64,
    total_chunks: usize,
) -> Vec<TransferTask> {
    let mut tasks = Vec::new();
    
    for i in 0..total_chunks {
        let offset = (i as u64) * (CHUNK_SIZE as u64);
        let remaining = file_size - offset;
        let chunk_size = std::cmp::min(CHUNK_SIZE as u64, remaining) as usize;
        
        tasks.push(TransferTask {
            chunk_info: ChunkInfo {
                index: i,
                offset,
                size: chunk_size,
            },
            file_path: file_path.to_string(),
            server_addr: server_addr.to_string(),
            task_id: task_id.to_string(),
        });
    }
    
    tasks
}

// 👷 工作协程主函数
async fn worker_routine(
    worker_id: usize,
    receiver: Receiver<TransferTask>,
    progress: std::sync::Arc<std::sync::Mutex<ProgressBar>>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    while let Ok(task) = receiver.recv() {
        // 📤 发送分片数据（零拷贝优化版本）
        if let Err(e) = send_chunk(
            &task.server_addr,
            &task.task_id,
            task.chunk_info.index,
            &task.file_path,
            task.chunk_info.offset,
            task.chunk_info.size
        ).await {
            eprintln!("❌ 工作协程{}发送分片{}失败: {}", worker_id, task.chunk_info.index, e);
            return Err(e.into());
        }
        
        // 更新进度条
        {
            let progress_lock = progress.lock().unwrap();
            progress_lock.inc(1);
        }
    }
    
    Ok(())
}

// 📤 发送单个分片数据（零拷贝优化版本）
async fn send_chunk(
    server_addr: &str,
    task_id: &str,
    chunk_index: usize,
    file_path: &str,
    offset: u64,
    size: usize,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    // 📡 建立TCP连接
    let mut stream = TcpStream::connect(server_addr).await?;
    
    // 📦 构造分片消息头: "CHUNK|任务ID|索引|大小"
    let header = format!("CHUNK|{}|{}|{}\n", task_id, chunk_index, size);
    
    // 📤 发送消息头
    stream.write_all(header.as_bytes()).await?;
    
    // 🔄 零拷贝传输：直接从文件读取并发送到网络连接
    let mut file = File::open(file_path)?;
    file.seek(SeekFrom::Start(offset))?;
    
    // 读取分片数据（避免一次性加载整个分片到内存）
    let mut buffer = vec![0u8; std::cmp::min(size, 64 * 1024)]; // 使用64KB缓冲区
    let mut remaining = size;
    while remaining > 0 {
        let to_read = std::cmp::min(buffer.len(), remaining);
        let read = file.read(&mut buffer[..to_read])?;
        if read == 0 {
            break;
        }
        stream.write_all(&buffer[..read]).await?;
        remaining -= read;
    }
    
    // 📥 等待服务器确认
    let mut buffer = [0; 1024];
    let n = stream.read(&mut buffer).await?;
    let response = String::from_utf8_lossy(&buffer[..n]);
    
    // ✅ 检查服务器响应
    if response.trim() != "CHUNK_RECEIVED" {
        return Err(format!("服务器未确认接收分片{}", chunk_index).into());
    }
    
    Ok(())
}